{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "from typing import *\n",
    "\n",
    "import metadsl as m\n",
    "import metadsl_core as mc\n",
    "import metadsl_visualize.visualize\n",
    "\n",
    "metadsl_visualize.visualize.monkeypatch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Array(m.Expression):\n",
    "    @m.expression\n",
    "    @classmethod\n",
    "    def create(cls) -> Array:\n",
    "        ...\n",
    "\n",
    "    @m.expression\n",
    "    def __getitem__(self, idx: Union[int, slice]) -> Array:\n",
    "        ...\n",
    "    \n",
    "    @m.expression\n",
    "    def index_int(self, i: mc.Integer) -> Array:\n",
    "        ...\n",
    "\n",
    "    @m.expression\n",
    "    def index_slice(self, start: mc.Integer, stop: mc.Integer, step: mc.Integer) -> Array:\n",
    "        ...\n",
    "\n",
    "    @m.expression\n",
    "    @property\n",
    "    def get_index(self) -> Union[slice, int]:\n",
    "        ...\n",
    "\n",
    "#     @m.expression\n",
    "#     @property\n",
    "#     def length(self) -> mc.Integer:\n",
    "#         ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "        IPython.get_ipython()\n",
    "        and IPython.get_ipython().__class__.__name__ == \"ZMQInteractiveShell\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ZMQInteractiveShell'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IPython.get_ipython().__class__.__name__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "@mc.register\n",
    "@m.rule\n",
    "def _getitem_typed(a: Array, i: int, s: slice):\n",
    "    yield a[i], a.index_int(mc.Integer.from_int(i))\n",
    "    yield a[s], lambda: a.index_slice(\n",
    "        mc.Integer.from_int(s.start),\n",
    "        mc.Integer.from_int(s.stop),\n",
    "        mc.Integer.from_int(s.step if s.step is not None else 1)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "@mc.register\n",
    "@m.rule\n",
    "def _slice_of_slice(\n",
    "    a: Array,\n",
    "\n",
    "    i: mc.Integer,\n",
    "    j: mc.Integer,\n",
    "    k: mc.Integer,\n",
    "\n",
    "    l: mc.Integer,\n",
    "    m: mc.Integer,\n",
    "    n: mc.Integer\n",
    "):\n",
    "    \"\"\"    \n",
    "    a[i:j:k][l:m:n] == a[i + j, i + m, n * k]\n",
    "    \"\"\"\n",
    "    return (\n",
    "        a.index_slice(i, j, k).index_slice(l, m, n),\n",
    "        a.index_slice(i + j, i + m, n * k)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "@mc.register\n",
    "@m.rule\n",
    "def _get_index(i: int, j: int, k: int):\n",
    "    yield Array.create().index_int(mc.Integer.from_int(i)).get_index, i\n",
    "    yield Array.create().index_slice(\n",
    "        mc.Integer.from_int(i),\n",
    "        mc.Integer.from_int(j),\n",
    "        mc.Integer.from_int(k),\n",
    "    ).get_index, slice(i, j, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/x.typez+json": {
       "nodes": [
        {
         "function": "Array.create",
         "id": "-308666362331542466"
        },
        {
         "id": "-3312730341382343964",
         "repr": "slice(1, 2, 2)",
         "type": "slice"
        },
        {
         "id": "-6362854663408583544",
         "repr": "slice(2, 3, 4)",
         "type": "slice"
        },
        {
         "args": [
          "-308666362331542466",
          "-3312730341382343964"
         ],
         "function": "Array.__getitem__",
         "id": "70989689021931579"
        },
        {
         "args": [
          "70989689021931579",
          "-6362854663408583544"
         ],
         "function": "Array.__getitem__",
         "id": "4138529908789473228"
        },
        {
         "args": [
          "4138529908789473228"
         ],
         "function": "Array.get_index",
         "id": "1443599856298358970"
        },
        {
         "id": "-1627799018779946948",
         "repr": "1",
         "type": "int"
        },
        {
         "id": "-1398281825510612154",
         "repr": "2",
         "type": "int"
        },
        {
         "args": [
          "-1627799018779946948"
         ],
         "function": "Integer.from_int",
         "id": "-230806222744429328"
        },
        {
         "args": [
          "-1398281825510612154"
         ],
         "function": "Integer.from_int",
         "id": "-6782319812949040250"
        },
        {
         "args": [
          "-308666362331542466",
          "-230806222744429328",
          "-6782319812949040250",
          "-6782319812949040250"
         ],
         "function": "Array.index_slice",
         "id": "-5747151165324484849"
        },
        {
         "args": [
          "-5747151165324484849",
          "-6362854663408583544"
         ],
         "function": "Array.__getitem__",
         "id": "-7376569648619854878"
        },
        {
         "args": [
          "-7376569648619854878"
         ],
         "function": "Array.get_index",
         "id": "-965015933505192850"
        },
        {
         "id": "2838897809788406090",
         "repr": "3",
         "type": "int"
        },
        {
         "id": "-3977614285577115941",
         "repr": "4",
         "type": "int"
        },
        {
         "args": [
          "2838897809788406090"
         ],
         "function": "Integer.from_int",
         "id": "-1259140478248144422"
        },
        {
         "args": [
          "-3977614285577115941"
         ],
         "function": "Integer.from_int",
         "id": "51497970396965384"
        },
        {
         "args": [
          "-5747151165324484849",
          "-6782319812949040250",
          "-1259140478248144422",
          "51497970396965384"
         ],
         "function": "Array.index_slice",
         "id": "-7880616000541017629"
        },
        {
         "args": [
          "-7880616000541017629"
         ],
         "function": "Array.get_index",
         "id": "-5784112661192481920"
        },
        {
         "args": [
          "-230806222744429328",
          "-6782319812949040250"
         ],
         "function": "Integer.__add__",
         "id": "8936858748525213959"
        },
        {
         "args": [
          "-230806222744429328",
          "-1259140478248144422"
         ],
         "function": "Integer.__add__",
         "id": "8166219812847196347"
        },
        {
         "args": [
          "51497970396965384",
          "-6782319812949040250"
         ],
         "function": "Integer.__mul__",
         "id": "4535011296845780294"
        },
        {
         "args": [
          "-308666362331542466",
          "8936858748525213959",
          "8166219812847196347",
          "4535011296845780294"
         ],
         "function": "Array.index_slice",
         "id": "970118383109874211"
        },
        {
         "args": [
          "970118383109874211"
         ],
         "function": "Array.get_index",
         "id": "2926188748599561587"
        },
        {
         "args": [
          "-308666362331542466",
          "-1259140478248144422",
          "8166219812847196347",
          "4535011296845780294"
         ],
         "function": "Array.index_slice",
         "id": "5219602393492759973"
        },
        {
         "args": [
          "5219602393492759973"
         ],
         "function": "Array.get_index",
         "id": "-6197227224833376235"
        },
        {
         "args": [
          "-308666362331542466",
          "-1259140478248144422",
          "51497970396965384",
          "4535011296845780294"
         ],
         "function": "Array.index_slice",
         "id": "3765905494179165203"
        },
        {
         "args": [
          "3765905494179165203"
         ],
         "function": "Array.get_index",
         "id": "41239340595245058"
        },
        {
         "id": "-4225717808849107396",
         "repr": "8",
         "type": "int"
        },
        {
         "args": [
          "-4225717808849107396"
         ],
         "function": "Integer.from_int",
         "id": "3613702820069102233"
        },
        {
         "args": [
          "-308666362331542466",
          "-1259140478248144422",
          "51497970396965384",
          "3613702820069102233"
         ],
         "function": "Array.index_slice",
         "id": "4502770117544835491"
        },
        {
         "args": [
          "4502770117544835491"
         ],
         "function": "Array.get_index",
         "id": "3958936131959955806"
        },
        {
         "id": "-5058553140429415303",
         "repr": "slice(3, 4, 8)",
         "type": "slice"
        }
       ],
       "states": {
        "initial": "1443599856298358970",
        "states": [
         {
          "node": "-965015933505192850",
          "rule": "__main__._getitem_typed"
         },
         {
          "node": "-5784112661192481920",
          "rule": "__main__._getitem_typed"
         },
         {
          "node": "2926188748599561587",
          "rule": "__main__._slice_of_slice"
         },
         {
          "node": "-6197227224833376235",
          "rule": "metadsl_core.integer.integer_math"
         },
         {
          "node": "41239340595245058",
          "rule": "metadsl_core.integer.integer_math"
         },
         {
          "node": "3958936131959955806",
          "rule": "metadsl_core.integer.integer_math"
         },
         {
          "node": "-5058553140429415303",
          "rule": "__main__._get_index"
         },
         {
          "label": "core",
          "node": "-5058553140429415303",
          "rule": ""
         }
        ]
       }
      },
      "text/plain": [
       "Typez(definitions=None, nodes=[CallNode(id='-308666362331542466', function='Array.create', type_params=None, args=None, kwargs=None), PrimitiveNode(id='-3312730341382343964', type='slice', repr='slice(1, 2, 2)'), PrimitiveNode(id='-6362854663408583544', type='slice', repr='slice(2, 3, 4)'), CallNode(id='70989689021931579', function='Array.__getitem__', type_params=None, args=['-308666362331542466', '-3312730341382343964'], kwargs=None), CallNode(id='4138529908789473228', function='Array.__getitem__', type_params=None, args=['70989689021931579', '-6362854663408583544'], kwargs=None), CallNode(id='1443599856298358970', function='Array.get_index', type_params=None, args=['4138529908789473228'], kwargs=None), PrimitiveNode(id='-1627799018779946948', type='int', repr='1'), PrimitiveNode(id='-1398281825510612154', type='int', repr='2'), CallNode(id='-230806222744429328', function='Integer.from_int', type_params=None, args=['-1627799018779946948'], kwargs=None), CallNode(id='-6782319812949040250', function='Integer.from_int', type_params=None, args=['-1398281825510612154'], kwargs=None), CallNode(id='-5747151165324484849', function='Array.index_slice', type_params=None, args=['-308666362331542466', '-230806222744429328', '-6782319812949040250', '-6782319812949040250'], kwargs=None), CallNode(id='-7376569648619854878', function='Array.__getitem__', type_params=None, args=['-5747151165324484849', '-6362854663408583544'], kwargs=None), CallNode(id='-965015933505192850', function='Array.get_index', type_params=None, args=['-7376569648619854878'], kwargs=None), PrimitiveNode(id='2838897809788406090', type='int', repr='3'), PrimitiveNode(id='-3977614285577115941', type='int', repr='4'), CallNode(id='-1259140478248144422', function='Integer.from_int', type_params=None, args=['2838897809788406090'], kwargs=None), CallNode(id='51497970396965384', function='Integer.from_int', type_params=None, args=['-3977614285577115941'], kwargs=None), CallNode(id='-7880616000541017629', function='Array.index_slice', type_params=None, args=['-5747151165324484849', '-6782319812949040250', '-1259140478248144422', '51497970396965384'], kwargs=None), CallNode(id='-5784112661192481920', function='Array.get_index', type_params=None, args=['-7880616000541017629'], kwargs=None), CallNode(id='8936858748525213959', function='Integer.__add__', type_params=None, args=['-230806222744429328', '-6782319812949040250'], kwargs=None), CallNode(id='8166219812847196347', function='Integer.__add__', type_params=None, args=['-230806222744429328', '-1259140478248144422'], kwargs=None), CallNode(id='4535011296845780294', function='Integer.__mul__', type_params=None, args=['51497970396965384', '-6782319812949040250'], kwargs=None), CallNode(id='970118383109874211', function='Array.index_slice', type_params=None, args=['-308666362331542466', '8936858748525213959', '8166219812847196347', '4535011296845780294'], kwargs=None), CallNode(id='2926188748599561587', function='Array.get_index', type_params=None, args=['970118383109874211'], kwargs=None), CallNode(id='5219602393492759973', function='Array.index_slice', type_params=None, args=['-308666362331542466', '-1259140478248144422', '8166219812847196347', '4535011296845780294'], kwargs=None), CallNode(id='-6197227224833376235', function='Array.get_index', type_params=None, args=['5219602393492759973'], kwargs=None), CallNode(id='3765905494179165203', function='Array.index_slice', type_params=None, args=['-308666362331542466', '-1259140478248144422', '51497970396965384', '4535011296845780294'], kwargs=None), CallNode(id='41239340595245058', function='Array.get_index', type_params=None, args=['3765905494179165203'], kwargs=None), PrimitiveNode(id='-4225717808849107396', type='int', repr='8'), CallNode(id='3613702820069102233', function='Integer.from_int', type_params=None, args=['-4225717808849107396'], kwargs=None), CallNode(id='4502770117544835491', function='Array.index_slice', type_params=None, args=['-308666362331542466', '-1259140478248144422', '51497970396965384', '3613702820069102233'], kwargs=None), CallNode(id='3958936131959955806', function='Array.get_index', type_params=None, args=['4502770117544835491'], kwargs=None), PrimitiveNode(id='-5058553140429415303', type='slice', repr='slice(3, 4, 8)')], states=States(initial='1443599856298358970', states=[State(node='-965015933505192850', rule='__main__._getitem_typed', label=None), State(node='-5784112661192481920', rule='__main__._getitem_typed', label=None), State(node='2926188748599561587', rule='__main__._slice_of_slice', label=None), State(node='-6197227224833376235', rule='metadsl_core.integer.integer_math', label=None), State(node='41239340595245058', rule='metadsl_core.integer.integer_math', label=None), State(node='3958936131959955806', rule='metadsl_core.integer.integer_math', label=None), State(node='-5058553140429415303', rule='__main__._get_index', label=None), State(node='-5058553140429415303', rule='', label='core')]))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "slice(3, 4, 8)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Array.create()[1:2:2][2:3:4].get_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @mc.register\n",
    "# @m.rule\n",
    "# def _create_length(name: str, l: Optional[int], l_integer: mc.Integer):\n",
    "#     def result():\n",
    "#         if l is None:\n",
    "#             raise m.NoMatch()\n",
    "#         return Array.create_length(name, mc.Integer.from_int(l))\n",
    "\n",
    "#     yield Array.create(name, l), result,\n",
    "#     yield Array.create_length(name, l_integer).length, l_integer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
